from tabnanny import check
import numpy as np
import copy
import matplotlib.pyplot as plt
from matplotlib.patches import Circle 
import time 
from math import pi, e, pow

def normalize(vector):
    return vector / np.linalg.norm(vector)
def Gaussian_function(x = 0 ,sigma = 1):
    return 0.4/sigma*pow(e, -x**2/(2*sigma**2))

class Node:
    def __init__(self, pos = None, expansion = -1, locked = False, nearest_ob = None) :
        self.pos = pos          # X, Y numpy array
        self.locked = locked    # if true, forces have no effect on this node
        self.expansion = expansion
        self.nearest_ob = nearest_ob
               
class EBNode:
    def __init__(self, maxExpension = 100, minExpension = 0.01, k_inter = 0.1, k_extel = 1, 
                 circle = 1000, path_arr = None, obs_arr = None, ob_radius = 30, pause_sec = 1):
        self.maxExpension = maxExpension
        self.minExpension = minExpension
        self.k_inter = k_inter        # 1 / how much a froce effects that node
        self.k_extel = k_extel  # coefficient for the tension force
        self.circle = circle
        self.ob_radius = ob_radius
        self.path_arr = path_arr
        self.obs_arr = obs_arr
        self.bobBands = []
        self.delete_num = []
        self.pause = pause_sec
        
    def getExpension(self, pos, obs):
        mincost = self.maxExpension
        nearest = [-1]
        for i in range(len(obs)):
            temp = np.linalg.norm(np.array(obs[i]) - pos) -self.ob_radius
            if (temp < mincost):
                mincost = temp
                nearest = np.array(obs[i])
        if mincost < self.minExpension:
            # return 0, nearest
            return self.minExpension, nearest
        return mincost, nearest
    
    def getExpension2(self, pos, obs):
        first_time = True
        mincost = 0
        for i in range(len(obs)):
            if first_time:
                mincost = np.linalg.norm(np.array(obs[i]) - pos) -self.ob_radius
                first_time = False
            else:
                temp = np.linalg.norm(np.array(obs[i]) - pos) -self.ob_radius
                if temp <mincost:
                    mincost = temp
        return mincost
    
    def getTotalForce(self, prv_node, current_node, nxt_node):
        difference1 = prv_node.pos - current_node.pos
        difference2 = nxt_node.pos - current_node.pos
        distance1   = np.linalg.norm(difference1)
        distance2   = np.linalg.norm(difference2)
        tenseForceDir = difference1/distance1+difference2/distance2
        tensionForce = self.k_inter*tenseForceDir
            
        if current_node.nearest_ob[0]!= -1 :
            if np.linalg.norm(tenseForceDir) >0.01:                   #小于0.01说明拉成直线了
                difference3 = current_node.pos - current_node.nearest_ob
                cos_theta = np.dot(difference3, tenseForceDir) / np.linalg.norm(tenseForceDir)/ np.linalg.norm(difference3)
                # cos_theta表示rep tense方向分量
                tensionForce = tensionForce + self.k_extel*tenseForceDir*cos_theta*Gaussian_function(x=current_node.expansion)
            else:
                difference3 = current_node.pos - current_node.nearest_ob
                distance3 = np.linalg.norm(difference3)
                tensionForce = tensionForce + self.k_extel*tenseForceDir*difference3/distance3*Gaussian_function(x=current_node.expansion)
            
        return tensionForce
    
    def suppressTangentialForces(self, node, forces):
        difference = node.prv.pos - node.nxt.pos
        scalar_fd = np.matmul(forces, difference)
        forces = forces - scalar_fd*difference
        return forces
    
    def applyForces(self, curren_node, force):
        curren_node.pos = curren_node.pos + curren_node.expansion*force
    
    def checkOverlap(self, check1, check2):
        if np.linalg.norm(self.bobBands[check1].pos - self.bobBands[check2].pos)<  \
                  (self.bobBands[check1].expansion+self.bobBands[check2].expansion):
            return True
        
    def Check(self, start_index, end_index):   
        if self.checkOverlap(start_index, end_index):
            if start_index+1<end_index:                   # 如果overlap且有其他值，则加入删除index
                for i in range(start_index+1,end_index):
                    self.delete_num.append([i, "del"])
            else:                                         # 如果overlap且中间没有其他的，直接返回了
                return                                    
        else:                                             # 没有overlap且中间还有值，开始递归
            if start_index+1<end_index:                   
                mid = (start_index + end_index) // 2
                self.Check(start_index, mid)
                self.Check(mid, end_index)
            else:                                          # 没有overlap且中间没有值，需要插值了
                self.delete_num.append([end_index, "insert"])
                   
    def removeAndFill(self):
        self.Check(0, len(self.bobBands)-1)
        self.delete_num.sort(reverse=True)                # 删除列表，从后往前删除
        for index in self.delete_num:
            if(index[1]=="del"):
                del(self.bobBands[index[0]])
            else:
                point = self.bobBands[index[0]-1].pos + (self.bobBands[index[0]].pos - self.bobBands[index[0]-1].pos)/2
                distance, near = self.getExpension(np.array(point), self.obs_arr)
                self.bobBands.insert(index[0],copy.deepcopy(Node(np.array(point), distance, locked = False, nearest_ob = near)))        
        self.delete_num.clear()
    
    def refreshBands(self, obs_arr):
        for bob in self.bobBands:
            distance, near = self.getExpension(bob.pos, obs_arr)
            bob.expansion =  distance
            bob.nearest_ob =  near
        self.removeAndFill()
    
    def bands_origin(self, path_arr, obs_arr):
       # 最开始的bobs
        distance, near = self.getExpension(np.array(path_arr[0]), obs_arr)
        self.bobBands.append(copy.deepcopy(Node(np.array(path_arr[0]), distance, True)))
        path_arr = np.delete(path_arr, 0, axis=0)
        for pos in path_arr:
            distance, near = self.getExpension(np.array(pos), obs_arr)
            self.bobBands.append(copy.deepcopy(Node(np.array(pos), distance, locked = False, nearest_ob = near)))
        self.bobBands[-1].locked = True             # 最后一个上锁
        
    def run(self):
        dis = []
        # rospy.init_node('te')
        fig, ax = plt.subplots()

        deepc = copy.deepcopy(self.obs_arr)
    # 生成bob band 
        self.bands_origin(self.path_arr, self.obs_arr)
        # self.refreshBands(self.obs_arr,0,len(self.bobBands))                             #  《-----------------

        real_path = [self.bobBands[0].pos]       # 每次完整更新一次，走一步
        path_toal = []
        path_toal.append(copy.deepcopy(real_path))
        edge_max = len(self.bobBands)

    # 边走边更新bob
        # for counter in range(edge_max):
        counter = -1
        # rate = rospy.Rate(10)
        while True:
            counter = 0
            start = time.time()

            edge_max = len(self.bobBands)
            # 设置一次只更新200个bob
            end_edge = counter+100
            if end_edge >= edge_max:
                end_edge = edge_max-2
                
            start_edge = 0
            if start_edge >= edge_max:
                break
            # 开始更新

            for i in range(self.circle):
                edge_max = len(self.bobBands)
                # 设置一次只更新200个bob
                end_edge = counter+100
                if end_edge >= edge_max:
                    end_edge = edge_max-2
                    
                start_edge = counter
                if start_edge >= edge_max :
                    break
            
                totForce_arr = []
                for j in range(start_edge,end_edge):
                    if not self.bobBands[j].locked:
                        totForce = self.getTotalForce(self.bobBands[j-1],self.bobBands[j],self.bobBands[j+1])                        #  《-----------------
                        totForce_arr.append(copy.deepcopy(totForce))
                    else:
                        totForce_arr.append(np.array([0,0]))   
                
                # # # 统一修改bobs  
                for m in range(len(totForce_arr)):
                    ind = start_edge +m
                    node = self.bobBands[ind]
                    if not node.locked:
                        self.applyForces(node, totForce_arr[m])             #  《-----------------
                self.refreshBands(self.obs_arr)                             #  《-----------------

            if start_edge+1<len(self.bobBands):
                move = self.bobBands[1].pos-self.bobBands[0].pos
                real_path.append(self.bobBands[0].pos+move/np.linalg.norm(move)/300.0)       # 每次完整更新一次，走一步
            # print(np.linalg.norm(self.bobBands[1].pos-self.bobBands[0].pos)/6.0)
            path_toal.append(copy.deepcopy(real_path))
            # print(np.linalg.norm(move))

    
            self.obs_arr[0,0] = self.obs_arr[0,0] - 0.001
            # self.obs_arr[0,1] = self.obs_arr[0,1] - 0.001
            # print(self.obs_arr)
            
            del(self.bobBands[0])
            distance, near = self.getExpension(np.array(real_path[-1]), self.obs_arr)
            dis.append(self.getExpension2(np.array(real_path[-1]), self.obs_arr)+0.040)
            self.bobBands.insert(0,copy.deepcopy(Node(np.array(real_path[-1]), distance, locked = True, nearest_ob = near)))  
            # print(real_path[-1])
            
            endj = self.bobBands[-1].pos-self.bobBands[0].pos
            if(np.linalg.norm(endj)<0.01):
                break
            
            
            
        return path_toal[-1],dis
        # self.drawPath(path, deepc)
    ########################3 take the path ###################################
                      
    def drawPath(self, path, obs_arr):
        fig = plt.figure()
        ax = fig.gca(projection='3d')
        for i in range(len(path)):
            pathi= np.array(path[i])
            ax.plot(pathi[:, 0], pathi[:, 1], pathi[:,2], color = "b", label='parametric curve')
            plt.pause(self.pause) 
            plt.cla()
# if __name__ == "__main__":
#     path_arr = [
#             [90.0, 90.0],
#             [131.0, 100.0],
#             [180.0,  80.0],
#             [200.0, 150.0],
#             [250.0, 190.0],
#             [200.0, 250.0]]    

#     obs_arr = [[150.0, 130.0],
#                 [170.0, 170.0]]
#     NHEBNode = EBNode(maxExpension = 100, minExpension = 0.01, k_inter = 0.1, k_extel = 1, 
#                  circle = 100, path_arr=path_arr,obs_arr=obs_arr, ob_radius=30)
#     NHEBNode.run()
    